#include "DensityOfStates.h"

const string DensityOfStatesWrapper::biassymbol = "Bias";
const string DensityOfStatesWrapper::wavevectorsymbol = "WaveVector";
const string DensityOfStatesWrapper::etasymbol = "Eta";
const string DensityOfStatesWrapper::nsamplessymbol = "Samples";
const string DensityOfStatesWrapper::npolysymbol = "ChebyshevStates";
const string DensityOfStatesWrapper::nrandsymbol = "RandomStates";
const string DensityOfStatesWrapper::nstatessymbol = "LanczosStates";
const string DensityOfStatesWrapper::methodsymbol = "Method";
const string DensityOfStatesWrapper::DMsymbol = "DirectMethod";
const string DensityOfStatesWrapper::KPMsymbol = "KernelPolynomialMethod";
const string DensityOfStatesWrapper::HMsymbol = "HaydockMethod";
const string DensityOfStatesWrapper::progresssymbol = "ProgressSymbol";

tbb::mutex DensityOfStatesWrapper::lock = tbb::mutex();

void DensityOfStatesWrapper::getDoS(int id)
{
	// The energy range is supplied first!
	numeric::range energy_range;
	if(!MLGetValue(stdlink,&energy_range))
		return;
	
	// The result is then allocated.
	mat result(2,energy_range.size());
	result.row(0).fill(energy_range.step);
	result(0,0) = energy_range.begin;
	result.row(0) = arma::cumsum(result.row(0));
	
	// The options are then supplied.
	rule::list opts;
	if(!MLGetValue(stdlink,&opts))
		return;

	// The method is optained!
	// Now we set the options!
	const size_t rulesize = opts.size();
	for(size_t r = 0; r < rulesize; r++)
	{
		// Rules from Mathematica always have symbols first.
		if(!opts[r].first.isSymbolType())
			continue;
		
		// The variable is pulled out.
		symbol sym = opts[r].first.toSymbol();
		variant var = opts[r].second;
		
		if(sym.name() == methodsymbol && var.isStringType())
		{
			if(var.toString() == DMsymbol)
			{
				DensityOfStates dos = getDOSWithParameters(opts);
				dos.E = result.row(0).t();
				result.row(1) = dos(id).DoS;
				break;
			}
			else if(var.toString() == KPMsymbol)
			{
				KernelPolynomialMethod kpm = getKPMWithParameters(opts);
				kpm.x.set_size(energy_range.size());
				kpm.x.fill(0.);
				kpm.x.set_real(result.row(0).t());
				result.row(1) = kpm.calcDoS(id).DoS.t();
				break;
			}
			else if(var.toString() == HMsymbol)
			{
				HaydockMethod hm = getHMWithParameters(opts);
				hm.E = result.row(0).t();
				result.row(1) = -arma::imag(hm(id).G)/arma::datum::pi;
				break;
			}
		}
		else if(sym.name() == methodsymbol && var.isSymbolType() && var.toSymbol().getType() == symbol::predefined && var.toSymbol().reference() == symbol::automatic)
		{
			if(System::Std.getUnit(id)->getAtoms().size() > 1000)
			{
				HaydockMethod hm = getHMWithParameters(opts);
				hm.E = result.row(0).t();
				result.row(1) = -arma::imag(hm(id).G)/arma::datum::pi;
				break;
			}
			else
			{
				DensityOfStates dos = getDOSWithParameters(opts);
				dos.E = result.row(0).t();
				result.row(1) = dos(id).DoS;
				break;
			}
		}
	}
	
	MLPutValue(stdlink,result);
}

DensityOfStates DensityOfStatesWrapper::getDOSWithParameters(const rule::list &opts)
{
	// The resulting functor is allocated!
	DensityOfStates dos;

	// Now we set the options!
	const size_t rulesize = opts.size();
	for(size_t r = 0; r < rulesize; r++)
	{
		// Rules from Mathematica always have symbols first.
		if(!opts[r].first.isSymbolType())
			continue;
		
		// The variable is pulled out.
		symbol sym = opts[r].first.toSymbol();
		variant var = opts[r].second;
		
		if(sym.name() == biassymbol)
		{
			if(var.isNumericType())
				dos.bias = var.toDouble();
		}
		else if(sym.name() == etasymbol)
		{
			if(var.isNumericType())
				dos.eta = var.toDouble();
			else if(var.isSymbolType() && var.toSymbol().reference() == symbol::none)
				dos.eta = 0;
		}
		else if(sym.name() == nsamplessymbol)
		{
			if(var.isNumericType())
				dos.N_samples = var.toSize();
		}
		else if(sym.name() == progresssymbol && var.isSymbolType() && var.toSymbol().getType() == symbol::userdefined)
		{
			string symstr = var.toSymbol().name();
			MLSetProgress(stdlink,symstr.c_str(),0.);
			auto updateprogress = [&,symstr] (const double &progress_step) {
				lock.lock();
				MLAddToProgress(stdlink, symstr.c_str(), progress_step);
				lock.unlock();
			};
			dos.updateProgress = updateprogress;
		}
	}

	return dos;
}
	
KernelPolynomialMethod DensityOfStatesWrapper::getKPMWithParameters(const rule::list &opts)
{
	// The resulting functor is allocated!
	KernelPolynomialMethod kpm;

	// Now we set the options!
	const size_t rulesize = opts.size();
	for(size_t r = 0; r < rulesize; r++)
	{
		// Rules from Mathematica always have symbols first.
		if(!opts[r].first.isSymbolType())
			continue;
		
		// The variable is pulled out.
		symbol sym = opts[r].first.toSymbol();
		variant var = opts[r].second;
		
		if(sym.name() == biassymbol)
		{
			if(var.isNumericType())
				kpm.bias = var.toDouble();
		}
		else if(sym.name() == wavevectorsymbol)
		{
			if(var.isNumericListType())
				kpm.k = var.toArmaDoubleVector3();
		}
		else if(sym.name() == etasymbol)
		{
			if(var.isNumericType())
				kpm.tol_bounds = var.toDouble();
			else if(var.isSymbolType() && var.toSymbol().reference() == symbol::none)
				kpm.tol_bounds = 0;
		}
		else if(sym.name() == nrandsymbol)
		{
			if(var.isNumericType())
				kpm.N_rand = var.toSize();
		}
		else if(sym.name() == npolysymbol || sym.name() == nsamplessymbol)
		{
			if(var.isNumericType())
				kpm.N_poly = var.toSize();
		}
		else if(sym.name() == nstatessymbol)
		{
			if(var.isNumericType())
				kpm.N_states = var.toSize();
		}
		else if(sym.name() == progresssymbol && var.isSymbolType() && var.toSymbol().getType() == symbol::userdefined)
		{
			string symstr = var.toSymbol().name();
			MLSetProgress(stdlink,symstr.c_str(),0.);
			auto updateprogress = [&,symstr] (const double &progress_step) {
				lock.lock();
				MLAddToProgress(stdlink, symstr.c_str(), progress_step);
				lock.unlock();
			};
			kpm.updateProgress = updateprogress;
		}
	}

	return kpm;
}
HaydockMethod DensityOfStatesWrapper::getHMWithParameters(const rule::list &opts)
{
	// The resulting functor is allocated!
	HaydockMethod hm;

	// Now we set the options!
	const size_t rulesize = opts.size();
	for(size_t r = 0; r < rulesize; r++)
	{
		// Rules from Mathematica always have symbols first.
		if(!opts[r].first.isSymbolType())
			continue;
		
		// The variable is pulled out.
		symbol sym = opts[r].first.toSymbol();
		variant var = opts[r].second;
		
		if(sym.name() == biassymbol)
		{
			if(var.isNumericType())
				hm.bias = var.toDouble();
		}
		else if(sym.name() == wavevectorsymbol)
		{
			if(var.isNumericListType())
				hm.k = var.toArmaDoubleVector3();
		}
		else if(sym.name() == etasymbol)
		{
			if(var.isNumericType())
				hm.eta = var.toDouble();
			else if(var.isSymbolType() && var.toSymbol().reference() == symbol::none)
				hm.eta = 0;
		}
		else if(sym.name() == nrandsymbol)
		{
			if(var.isNumericType())
				hm.N_rand = var.toSize();
		}
		else if(sym.name() == nstatessymbol || sym.name() == nsamplessymbol)
		{
			if(var.isNumericType())
				hm.N_states = var.toSize();
		}
		else if(sym.name() == progresssymbol && var.isSymbolType() && var.toSymbol().getType() == symbol::userdefined)
		{
			string symstr = var.toSymbol().name();
			MLSetProgress(stdlink,symstr.c_str(),0.);

			auto updateprogress = [&,symstr] (const double &progress_step) {
				lock.lock();
				MLAddToProgress(stdlink, symstr.c_str(), progress_step);
				lock.unlock();
			};
			hm.updateProgress = updateprogress;
		}
	}

	return hm;
}

#include <sstream>
#include <iomanip>


bool DensityOfStatesWrapper::MLSetProgress(MLINK mlp, const char* var, double val) {
	
	std::stringstream s;

	s << "Progress[" << var << "] ^= " << std::fixed << std::setprecision(9) << val << ";";
	
	bool result = MLEvaluateString(stdlink, (char*)s.str().c_str());

	if(! result)
	{
		MLReturnError("GrapheneTools","error","When processing a value update the program failed to evaluate update.");
		return false;
	}

	return true;
}

bool DensityOfStatesWrapper::MLAddToProgress(MLINK mlp, const char* var, double val) {
	
	std::stringstream s;

	s << "Progress[" << var << "] ^= " << "Progress[" << var << "] + " << std::fixed << std::setprecision(9) << val << ";";

	if(! MLEvaluateString(stdlink, (char*)s.str().c_str()))
	{
		MLReturnError("GrapheneTools","error","When processing a value update the program failed to evaluate update.");
		return false;
	}

	return true;
}